import networkx as nx
import numpy as np
import matplotlib.pyplot as plt


class EpidemicModel:

    def __init__(self,
                 contact_networks: nx.Graph = None,
                 n=1000,
                 beta=0.05,
                 beta_decay=0.9,
                 infective_period=7,
                 protection_period=30,
                 num_worker=10) -> None:
        if contact_networks is not None:
            self.contact_networks = contact_networks
        else:
            import web_page_generate
            # self.contact_networks = nx.newman_watts_strogatz_graph(n, 10, 0.05)
            # self.contact_networks = web_page_generate.web_page_with_random_edge(
            #     n, 7, 3, 0.7)
            self.contact_networks = web_page_generate.web_page(n, 1, 0.7)
        if self.contact_networks.is_directed():
            self.contact_networks = self.contact_networks.to_undirected()

        self.beta = beta
        self.beta_decay = beta_decay
        self.infective_period = infective_period
        self.protection_period = protection_period
        self.num_worker = num_worker
        for node, attributes in self.contact_networks.nodes.data():
            attributes['state'] = 'S'
            attributes['beta'] = self.beta

        self.number_of_state = {
            'S': self.contact_networks.number_of_nodes(),
            'I': 0,
            'R': 0,
            'Isolated': 0
        }

        self.color_dict = {'S': 'orange', 'I': 'r', 'R': 'g', 'isolated': 'm'}

    def random_infect(self, num):
        """
        随机初始化感染点 num个
        """
        infected_node = np.random.choice(self.contact_networks.nodes(),
                                         size=num,
                                         replace=False)
        for node in infected_node:
            attributes = self.contact_networks.nodes[node]
            attributes['state'] = 'I'
            attributes['period'] = self.infective_period

        self.number_of_state['S'] -= num
        self.number_of_state['I'] += num

    def node_spread(self, node):
        """
        传染节点传播
        """
        # 记录新被感染的节点
        infected_node = set()
        # if self.contact_networks.nodes['state'] == 'I':
        for nbr in self.contact_networks.neighbors(node):
            attributes = self.contact_networks.nodes[nbr]
            if attributes['state'] == 'S' and np.random.random(
            ) < attributes['beta']:
                attributes['state'] = 'I'
                attributes['period'] = self.infective_period

                infected_node.add(nbr)

                self.number_of_state['S'] -= 1
                self.number_of_state['I'] += 1
        return infected_node

    def update(self):
        """
        update the network
        """
        new_infected_node = set()
        for node, attributes in self.contact_networks.nodes.data():
            if attributes['state'] == 'I' and node not in new_infected_node:
                # 当没有被隔离时
                if not attributes.get('isolated', False):
                    infected_node = self.node_spread(node)
                    new_infected_node.update(infected_node)
                attributes['period'] -= 1
                if attributes['period'] <= 0:
                    attributes['state'] = 'R'
                    attributes['period'] = self.protection_period

                    # 提高免疫力 降低感染概率
                    attributes['beta'] *= self.beta_decay

                    if attributes.get('isolated', False):
                        # 解除隔离
                        attributes['isolated'] = False
                        self.number_of_state['Isolated'] -= 1

                    self.number_of_state['I'] -= 1
                    self.number_of_state['R'] += 1
                elif not attributes.get('isolated', False):
                    new_infected_node.add(node)
            elif attributes['state'] == 'R':
                attributes['period'] -= 1
                if attributes['period'] <= 0:
                    attributes['state'] = 'S'

                    self.number_of_state['R'] -= 1
                    self.number_of_state['S'] += 1

        # 流调隔离
        if self.num_worker > 0 and len(new_infected_node) > 0:
            isolated_node = np.random.choice(list(new_infected_node),
                                             size=min(self.num_worker,
                                                      len(new_infected_node)),
                                             replace=False)
            for node in isolated_node:
                attributes = self.contact_networks.nodes[node]
                # if not attributes.get('isolated', False):
                attributes['isolated'] = True
                self.number_of_state['Isolated'] += 1

    def get_num_of_each_state(self):
        return self.number_of_state

    def get_node_color(self):
        return [
            self.color_dict['isolated' if attributes.
                            get('isolated', False) else attributes['state']]
            for _, attributes in self.contact_networks.nodes.data()
        ]


if __name__ == "__main__":
    total_time = 365 * 2
    g = nx.newman_watts_strogatz_graph(1000, 8, 0.05)
    # protection_period == 8时会震荡
    model = EpidemicModel(contact_networks=g,
                          n=1000,
                          beta=0.05,
                          beta_decay=1,
                          infective_period=14,
                          protection_period=14,
                          num_worker=100)
    # S = [model.contact_networks.number_of_nodes()]
    # I = [0]
    # R = [0]
    # Isolated = [0]
    S = np.zeros(total_time + 1)
    I = np.zeros(total_time + 1)
    R = np.zeros(total_time + 1)
    Isolated = np.zeros(total_time + 1)
    S[0] = model.contact_networks.number_of_nodes()
    I[0] = 0
    R[0] = 0
    Isolated[0] = 0

    # fig, axes = plt.subplots(2, 1, figsize=(5,10))
    # ax0 = axes[0]
    # ax1 = axes[1]
    # fig1 = plt.figure("contact network", figsize=(10,10))
    # ax0 = fig1.add_subplot()
    # fig2 = plt.figure("statistics", figsize=(5,5))
    # ax1 = fig2.add_subplot()
    # ax1.set_title("number of people")
    fig = plt.figure("epidemic model", figsize=(18, 8))
    ax0 = plt.subplot2grid((2, 3), (0, 0), rowspan=2, colspan=2, fig=fig)
    ax1 = plt.subplot2grid((2, 3), (0, 2), rowspan=2, fig=fig)
    ax0.set_title("contact network")
    ax1.set_title("statistics")
    ax1.set_xlabel("time/day")
    ax1.set_ylabel("number")

    pos = nx.spring_layout(model.contact_networks)

    for time in range(1, total_time + 1):
        if time == 1:
            model.random_infect(50)
        else:
            model.update()
        num_of_state = model.get_num_of_each_state()

        # S.append(num_of_state['S'])
        # I.append(num_of_state['I'])
        # R.append(num_of_state['R'])
        S[time] = num_of_state['S']
        I[time] = num_of_state['I']
        R[time] = num_of_state['R']
        Isolated[time] = num_of_state['Isolated']

        ax0.cla()
        ax1.cla()
        ax0.set_title(f"contact networks(day {time})")
        # ax1.set_xlim(0)
        nx.draw(model.contact_networks,
                pos=pos,
                ax=ax0,
                node_size=25,
                node_color=model.get_node_color(),
                edge_color='#D8D8D8',
                font_color='white')
        ax1.plot(S[:time + 1], c='orange', label="S")
        ax1.plot(I[:time + 1], c='r', label="I")
        ax1.plot(R[:time + 1], c='g', label="R")
        ax1.plot(Isolated[:time + 1], 'm--', label="Isolated")
        # ax1.plot(S, c='orange', label="S")
        # ax1.plot(I, c='r', label="I")
        # ax1.plot(R, c='g', label="R")
        # ax1.plot(Isolated, 'm--.', label="Isolated")
        fig.legend()
        plt.pause(0.001)
        print(
            f"S: {S[time]}    I: {I[time]}    R: {R[time]}    Isolated: {Isolated[time]}"
        )

        if I[time] <= 0:
            print("疫情已结束")
            print(
                f"S: {S[time]}    I: {I[time]}    R: {R[time]}    Isolated: {Isolated[time]}"
            )
            break
    input("finished. input any to exit")
